import fitz
import os
from PIL import Image
filename = "../demo.pdf"
search_term = "OVERCURRENT"
# 官网: https://pymupdf.readthedocs.io/en/latest/

# http://blog.yixao.net/share/12997.html
def save_png_set_pixmap_clip_by_search_keyword(filename, key_word):
    # 解决清晰度问题 resolution 分辨率
    zoom_x = 3.0  # horizontal zoom
    zoom_y = 3.0  # vertical zoom
    mat = fitz.Matrix(zoom_x, zoom_y)
    left_top_x_position = 20        # 左上x坐标
    left_top_y_gatter_position = 0        # 左上y坐标向上留空白的gatter(读取关键词需要重新赋值)
    bottom_right_x_gatter = 60
    bottom_right_y_gatter = 60
    pdf_document = fitz.open(filename)

    # 创建一个新的pdf, 并新建一个page
    out_pdf = fitz.open()
    # out_pdf_page = out_pdf.new_page()
    # insert_text_point = fitz.Point(20, 72)
    # insert_text = 'hello world'
    # out_pdf_page.insert_text(insert_text_point, insert_text, fontname="helv", fontsize=11, rotate=0)
    # out_pdf.save('insert_text.pdf')

    for current_page in range(len(pdf_document)):
        page = pdf_document.loadPage(current_page)
        result_page = page.searchFor(key_word)
        if result_page:
            # 相关关键词, 整页保存为pdf
            out_pdf.insert_pdf(pdf_document, from_page=page.number, to_page=page.number)
            out_pdf.save("out-%s.pdf" % key_word)

            # 获取pdf每页的像素宽高(0.0, 0.0, 792.0, 792.0)
            rect = page.rect
            page_bottom_right_position = rect.br

            # 剪裁的左上角的 (x, y) => x是相对固定值, 距离左边留有空余, y是关键词在文档中的位置向上移动gatter个像素
            left_top_y_position = result_page[0][1] - left_top_y_gatter_position
            bottom_right_x_position = page_bottom_right_position[0] - bottom_right_x_gatter
            bottom_right_y_position = page_bottom_right_position[1] - bottom_right_y_gatter

            # 最终剪裁区域的左上,右下的坐标
            clip_start = (left_top_x_position, left_top_y_position)
            clip_end = (bottom_right_x_position, bottom_right_y_position)

            # print("rect.tl, rect.br: ", rect.tl, rect.br)  # Point(0.0, 0.0) Point(612.0, 792.0)
            # 最终希望剪裁的区域
            clip = fitz.Rect(clip_start, clip_end)
            pix = page.get_pixmap(matrix=mat, clip=clip)
            pix.writePNG("page-%i-%s.png" % (page.number, key_word))

            # 生成图片pdf
            # w = 595
            # h = 842
            # out_pdf_page.insertImage(fitz.Rect(0, 0.85*h, w, h))
            # out_pdf.save('new-test.pdf')
            # pm = page.get_pixmap(matrix=mat, clip=clip)
            # pm.writeImage("%s-%s.jpg" % (key_word, page.number))

# 根据生成的多个img生成统一的pdf
# TODO: 现在只是做了每个截图存成一页的pdf, 效果不好, 需要计算每张图片的高度,再加上一个margin就可以组合成有限的pages
def combine_images_to_pdf(output_filename, base_path, endswith):
    out_pdf = fitz.open()
    image_list = get_image_path(base_path, endswith)
    for img_path in image_list:
        new_page = out_pdf.new_page()
        img = open(img_path, "rb").read()
        new_page.insertImage(fitz.Rect(30,30, 500, 800), stream=img)
    out_pdf.save(output_filename)


# def combine_images_to_pdf(output_filename, base_path, endswith):
#     out_pdf = fitz.open()
#     image_list = get_image_path(base_path, endswith)
#     for img_path in image_list:
#         new_page = out_pdf.new_page()
#         img = open(img_path, "rb").read()
#         temp_img = Image.open(img_path)
#         print(temp_img.size[1])
#         rect = fitz.Rect(0, 0, 600, temp_img.size[1])
#         start_point = rect.tl
#         end_point = rect.br
#         new_page.insertImage(fitz.Rect(start_point, end_point), stream=img)
#         start_point += (0, temp_img.size[1])
#         end_point += (0, temp_img.size[1])
#     out_pdf.save(output_filename)

# 给定路径, 以及查找的扩展名, 返回所有满足条件的文件的路径
def get_image_path(base_path, endswith):
    image_list = []
    collection = os.walk(base_path)
    for path, d, files in collection:
        for file in files:
            if file.endswith(endswith):
                image_list.append(os.path.join(path, file))
    return image_list


# 1. 生成拆分的图片
save_png_set_pixmap_clip_by_search_keyword(filename, search_term)
# 2. 将图片进行组合
combine_images_to_pdf("combine-pdf.pdf", "./", ".png")
